/**************************************************************************
* Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
*                                                                        *
* Author: The ALICE Off-line Project.                                    *
* Contributors are mentioned in the code where appropriate.              *
*                                                                        *
* Permission to use, copy, modify and distribute this software and its   *
* documentation strictly for non-commercial purposes is hereby granted   *
* without fee, provided that the above copyright notice appears in all   *
* copies and that both the copyright notice and this permission notice   *
* appear in the supporting documentation. The authors make no claims     *
* about the suitability of this software for any purpose. It is          *
* provided "as is" without express or implied warranty.                  *
**************************************************************************/
//
//  Xian-Guo Lu <lu@physi.uni-heidelberg.de>
//

#include "AliAnalysisManager.h"
#include "AliAODHandler.h"
#include "AliAODInputHandler.h"
#include "AliCentrality.h"
#include "AliESDtrackCuts.h"
#include "AliESDInputHandler.h"
#include "AliESDv0.h"
#include "AliESDv0KineCuts.h"
#include "AliTPCdEdxInfo.h"

#include "AliAnalysisTaskpkRaa.h"

AliAnalysisTaskpkRaa::AliAnalysisTaskpkRaa():
  AliAnalysisTaskSE("pkRaa")
  , fESDpid(0x0)
 
  , fEvent(0x0)
  , fStream(0x0)
  , fTree(0x0)
  , fList(0x0)
  , fCuts(0x0)

  , fCen(0x0)
  , fEta(0x0)
  , fTPCpidn(0x0)
  , fTPCnf(0x0)
  , fFilterBit(0x0)
  , fEntry(0x0)
  , fNsel(0x0)
{
  //
  //
  //

  DefineOutput(1, TTree::Class());
  DefineOutput(2, TList::Class());

  fStream = new TTreeStream("tree");
  fTree=fStream->GetTree();

  fList = new TList;
  fList->SetOwner(kTRUE);

  fCuts = new TH1I("cuts","", 20,0,20);
  fCen = new TH1D("cen","",110,-10,100);
  fEta = new TH1D("eta","",100,-1,1);
  fTPCpidn = new TH1I("tpcpidn","",200,0,200);
  fTPCnf = new TH1I("tpcnf","",200,0,200);
  fFilterBit = new TH1I("filterbit","",20,0,20);
  fEntry = new TH1I("entry","",2,0,2);
  fNsel = new TH1I("nsel","",100,0,100);

  fList->Add(fCuts);
  fList->Add(fNsel);
  fList->Add(fCen);
  fList->Add(fEta);
  fList->Add(fTPCpidn);
  fList->Add(fTPCnf);
  fList->Add(fFilterBit);
  fList->Add(fEntry);
}

AliAnalysisTaskpkRaa::~AliAnalysisTaskpkRaa()
{
  //
  // Destructor
  //

  /* can not be deleted in chain, they are in input handler!
  delete fEvent;
  */

  //delete fStream;
  delete fList;
}

void AliAnalysisTaskpkRaa::UserCreateOutputObjects()
{
  //
  //
  //

  PostData(1, fTree);
  PostData(2, fList);
}

void AliAnalysisTaskpkRaa::UserExec(Option_t *)
{
  //
  //Execute
  // 

  fCuts->Fill(0);
  if(!CutInput())
    return;

  fCuts->Fill(1);
  if(!CutEvent())
    return;

  fCuts->Fill(2);
  const Int_t ntrk0 = fEvent->GetNumberOfTracks();
  typetrack * trks[ntrk0];
  for(Int_t ii=0; ii<ntrk0; ii++){
    trks[ii] = 0x0;
  }
  Int_t nsel = 0;
  CutTrack(nsel, trks);
  fNsel->Fill(nsel);
  if(!nsel)
    return;

  fCuts->Fill(3);

  Float_t v0cen = -999;
  AliCentrality *evtCentrality = fEvent->GetCentrality();
  if(evtCentrality){
    v0cen = evtCentrality->GetCentralityPercentile("V0M");
    //apply centrality cut here to accelerate
    
    fCen->Fill(v0cen);
    if(v0cen<0 || v0cen>5){
      return;
    }
  }

  fCuts->Fill(4);

  //-------------
  Int_t pdgs[ntrk0];
  Int_t mother[ntrk0];
  Float_t cpa[ntrk0];
  Float_t dca[ntrk0];
  Float_t mass[ntrk0];
  Float_t arm0[ntrk0];
  Float_t arm1[ntrk0];
  Float_t psi[ntrk0];
  GetV0type(ntrk0, v0cen, pdgs, mother, cpa, dca, mass, arm0, arm1, psi);

  //-------------

  for(Int_t itrk=0; itrk<nsel ; itrk++){
    typetrack * trackptr = trks[itrk];

    Float_t mom = trackptr->P();
    Int_t charge = trackptr->Charge();
    Float_t pt = trackptr->Pt();
   
    Float_t eta = trackptr->Eta();
    fEta->Fill(eta);
    if(TMath::Abs(eta)<0.6 || TMath::Abs(eta)>0.8)
      continue;

    Int_t tpcncross = GetNCross(trackptr);
    AliTPCdEdxInfo *tpcinfo = trackptr->GetTPCdEdxInfo();
    Double32_t infosignal[4];
    Char_t infoncl[3];
    Char_t infonrows[3];
    if(tpcinfo){
      tpcinfo->GetTPCSignalRegionInfo(infosignal, infoncl, infonrows);
      tpcncross = infonrows[0]+infonrows[1]+infonrows[2];
    }

    Float_t tpcsig  = trackptr->GetTPCsignal();
    Float_t tpcip = trackptr->GetTPCmomentum();
    Int_t tpcnf = trackptr->GetTPCNclsF();
    Int_t tpcpidn = trackptr->GetTPCsignalN();
    fTPCpidn->Fill(tpcpidn);
    if(tpcpidn<80)
      continue;
    fTPCnf->Fill(tpcnf);
    if(tpcnf<10)
      continue;

    /*
    Float_t trdip = -999;
    for(Int_t ich=0; ich<6; ich++){
      trdip = trackptr->GetTRDmomentum(ich);
      if(trdip>0)
        break;
    }
    Float_t trdsig = trackptr->GetTRDsignal();
    Int_t trdnch  = trackptr->GetTRDNchamberdEdx();
    Int_t trdncls = trackptr->GetTRDNclusterdEdx();
    */

    Float_t tofpion = GetTOFNsigma(trackptr, AliPID::kPion);
    Float_t tofproton = GetTOFNsigma(trackptr, AliPID::kProton);
    Float_t tofkaon = GetTOFNsigma(trackptr, AliPID::kKaon);
    Float_t tofelectron = GetTOFNsigma(trackptr, AliPID::kElectron);

    Float_t hmpidsig = -999;
    Float_t hmpidprob[5]={-999,-999,-999,-999,-999};
    GetHMPID(trackptr, hmpidprob, hmpidsig);

    
    const Int_t trkid = trackptr->GetID();
    /*
    if(pdgs[trkid]!=fkBad && pdgs[trkid]!=fkNotUsed){
      printf("test v0 %5d: %5d %5d -- %5d %5d -- %.4f %.4f %.4f %.4f %.4f %.4f\n", trkid, ipar, tmpbit, pdgs[trkid], mother[trkid], cpa[trkid], dca[trkid], mass[trkid], arm0[trkid], arm1[trkid], psi[trkid]);
    }
    */

    static Int_t entry = 0;
    
    (*fStream)<<"tree"<<
      "v0cen="<<v0cen<<

      "mom="<<mom<<
      "charge="<<charge<<
      "pt="<<pt<<
      "eta="<<eta<<

      "tpcsig="<<tpcsig<<
      "tpcip="<<tpcip<<
      "tpcpidn="<<tpcpidn<<
      "tpcnf="<<tpcnf<<
      "tpcncross="<<tpcncross;

    if(tpcinfo){
      for(Int_t itmp=0; itmp<4; itmp++){
        (*fStream)<<"tree"<<
          Form("infosignal%d=",itmp)<<infosignal[itmp];
      }
      
      for(Int_t itmp=0; itmp<3; itmp++){
        (*fStream)<<"tree"<<
          Form("infoncl%d=",itmp)<<infoncl[itmp];
      }
      
      for(Int_t itmp=0; itmp<3; itmp++){
        (*fStream)<<"tree"<<
          Form("infonrows%d=",itmp)<<infonrows[itmp];
      }
    }

    if(IsESD()){
      (*fStream)<<"tree"<<
        /*
          "trdip="<<trdip<<
          "trdsig="<<trdsig<<
          "trdnch="<<trdnch<<
          "trdncls="<<trdncls<<
        */
        
        "tofpion="<<tofpion<<
        "tofproton="<<tofproton<<
        "tofkaon="<<tofkaon<<
        "tofelectron="<<tofelectron<<
        
        "hmpidElectron="<<hmpidprob[0]<<
        "hmpidPion="<<hmpidprob[2]<<
        "hmpidKaon="<<hmpidprob[3]<<
        "hmpidProton="<<hmpidprob[4]<<
        "hmpidsig="<<hmpidsig<<
        
        "pdg="<<pdgs[trkid]<<
        "mother="<<mother[trkid]<<
        "mass="<<mass[trkid]<<
        "arm0="<<arm0[trkid]<<
        "arm1="<<arm1[trkid];
    }

    (*fStream)<<"tree"<<
      "\n";
    
    fEntry->Fill(1);
    entry++;
  } 
}

Double_t AliAnalysisTaskpkRaa::GetNCross(TObject* obj)
{
  AliESDtrack * esdtrk = dynamic_cast<AliESDtrack*>(obj);
  if(esdtrk){
    return esdtrk->GetTPCCrossedRows();
  }
  else{
    return -999;
  }
}

Bool_t AliAnalysisTaskpkRaa::CutInput()
{
  //
  //general protection
  //
  fEvent = 0x0;

  AliVEventHandler * evtH = AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler();
  if(!evtH){
    printf("AliAnalysisTaskAODpkRaa::CutInput error No evtH\n"); exit(1);
  }

  if(evtH->InheritsFrom("AliESDInputHandler")){
    fEvent =  ((AliESDInputHandler*)evtH)->GetEvent();

    fESDpid = (AliESDInputHandler*) evtH->GetESDpid();  
    if(!fESDpid){
      printf("AliAnalysisTaskpkRaa::CutInput error no ESDpid!!\n"); exit(1);
    }
  }
  else if(evtH->InheritsFrom("AliAODInputHandler")){
    fEvent =  ((AliAODInputHandler*)evtH)->GetEvent();
  }
  else if(evtH->InheritsFrom("AliAODHandler")){
    fEvent = ((AliAODHandler*)evtH)->GetAOD();
  }
  else{
    printf("AliAnalysisTaskAODpkRaa::CutInput error No aodH type %s\n", evtH->GetDataType()); exit(1);
  }

  if(!fEvent){
    printf("AliAnalysisTaskAODpkRaa::CutInput error No AODEvent  AODInputHandler %d AODHandler %d\n", evtH->InheritsFrom("AliAODInputHandler"), evtH->InheritsFrom("AliAODHandler")); exit(1);
  }

  //in unit of kg
  const Float_t dm = fEvent->GetMagneticField();
  if( TMath::Abs(TMath::Abs(dm)- 5) > 0.1 ){
    printf("AliAnalysisTaskpkRaa::CutInput error strange Bfield! %f\n", dm); exit(1);
  }

  return kTRUE;
}

Bool_t AliAnalysisTaskpkRaa::CutEvent()
{
  const AliVertex *vertex = 0x0;
  
  if(IsESD()){
    vertex = fEvent->GetPrimaryVertexTracks();
    if(!vertex){
      return kFALSE;
    }
    if(vertex->GetNContributors()<1) {
      // SPD vertex
      vertex = fEvent->GetPrimaryVertexSPD();
      if(vertex->GetNContributors()<1) {
        // NO GOOD VERTEX, SKIP EVENT 
        return kFALSE;
      }
    }
  }
  else{
    vertex = fEvent->GetPrimaryVertex();
    if(!vertex){
      return kFALSE;
    }
    //3 types of vertex: Tracks, SPD, TPC
    //do not use TPC vertex
    const TString pvname(vertex->GetName());
    if(pvname=="TPCVertex"){
      return kFALSE;
    }
    
    if(vertex->GetNContributors()<1){
      return kFALSE;
    }
  }

  vertex = fEvent->GetPrimaryVertex();
  if(TMath::Abs(vertex->GetZ())>10 ){
    return kFALSE;
  }

  return kTRUE;
}

void AliAnalysisTaskpkRaa::CutTrack(Int_t &nsel,  typetrack *outtrk[])
{
  /*
  //____________________________________________________________________________
  //____________________________________________________________________________
  //10h_AODGSI /hera/alice/rromita/trunk_AOD/aod_AOD/AddTask_aod_ESDFilter_PbPb.C
  // standard cuts with tight DCA cut, using cluster cut instead of crossed rows (a la 2010 default)                                                          
  AliESDtrackCuts* esdTrackCutsH2Cluster = AliESDtrackCuts::GetStandardITSTPCTrackCuts2010(kTRUE, 0);
  // 1024, 1<< 10                                                              
  trackFilter->AddCuts(esdTrackCutsH2Cluster);

  //____________________________________________________________________________
  //____________________________________________________________________________
  //10h_AOD049
  //AOD049 	LHC10h 	  	28.5M 	v4-21-21-AN 	AddTaskESDFilter.C
  //http://alisoft.cern.ch/viewvc/tags/v4-21-21-AN/ANALYSIS/macros/AddTaskESDFilter.C?view=markup&root=AliRoot
  // standard cuts with tight DCA cut
  AliESDtrackCuts* esdTrackCutsH2 = AliESDtrackCuts::GetStandardITSTPCTrackCuts2010();
  
  // 32
  trackFilter->AddCuts(esdTrackCutsH2);

  //____________________________________________________________________________
  //____________________________________________________________________________
  //11h_AOD095
  //https://twiki.cern.ch/twiki/bin/viewauth/ALICE/PWGPPAODTrackCuts
  //AOD095 	LHC11h 	pass2 	94M 	v5-02-Rev-04 	AddTaskESDFilter.C 	
  //http://alisoft.cern.ch/viewvc/tags/v5-02-Rev-04/ANALYSIS/macros/AddTaskESDFilter.C?view=markup&root=AliRoot
  
  Bool_t AddTrackCutsLHC11h(AliAnalysisTaskESDfilter* esdfilter)
  // standard cuts with tight DCA cut, using cluster cut instead of crossed rows (a la 2010 default)
  AliESDtrackCuts* esdTrackCutsH2Cluster = AliESDtrackCuts::GetStandardITSTPCTrackCuts2011(kTRUE, 0);
  // 1024, 1<< 10 // tight DCA cuts
  trackFilter->AddCuts(esdTrackCutsH2Cluster);
  //____________________________________________________________________________
  //____________________________________________________________________________
  */

  //
  //CutTrack
  //
  
  const Int_t run = fEvent->GetRunNumber();
  //checked on E-log ->
  //139846 10h
  //139847 11a
  //170718 11h
  //170719 12a
  //<-

  if(IsESD()){
    AliESDtrackCuts * esdTrackCuts = 0x0;
    
    if(run<=139846){//2010
      esdTrackCuts = AliESDtrackCuts::GetStandardITSTPCTrackCuts2010(kTRUE, 0);
    }
    else if(run >=139847 && run<=170718){//2011
      esdTrackCuts = AliESDtrackCuts::GetStandardITSTPCTrackCuts2011(kTRUE, 0);
    }
    else{
      printf("AliAnalysisTaskpkRaa::CutTrack wrong run!! %d\n", run); exit(1);
    }
    
    nsel = 0;
    for(Int_t ii=0; ii<fEvent->GetNumberOfTracks(); ii++){
      typetrack *trackptr = fEvent->GetTrack(ii);
      if(esdTrackCuts->AcceptTrack(trackptr)){
        outtrk[nsel]=trackptr;
        //test 135 -- 135 -3141593  ==> so GetID is the correct one
        //printf("test %d -- %d %d\n", ii, outtrk[nsel]->GetID(), outtrk[nsel]->GetLabel());
        nsel++;
      }
    }
    delete esdTrackCuts;  
  }
  else{//AOD

    Int_t nshift = 10;
    if(run<=139846){//2010
      nshift = 5; //10h_AOD049
    }
    else if(run >=139847 && run<=170718){//2011
      nshift = 10; //11h_AOD095
    }
    else{
      printf("AliAnalysisTaskpkRaa::CutTrack wrong run!! %d\n", run); exit(1);
    }
    
    fFilterBit->Fill(nshift);
    
    const Int_t ibit = BIT(nshift);
    
    nsel = 0;
    for(Int_t ii=0; ii<fEvent->GetNumberOfTracks(); ii++){
      typetrack *trackptr = fEvent->GetTrack(ii);
      if(trackptr->TestFilterBit(ibit)){
        outtrk[nsel]=trackptr;
        nsel++;
      }
    }
  }
}


Float_t AliAnalysisTaskpkRaa::GetTOFNsigma(const typetrack *trk, const AliPID::EParticleType itype)
{
  //itype: AliPID::kElectron, AliPID::kPion, AliPID::kKaon, AliPID::kProton
  //
  //

  if(!IsESD()){
    return -999;
  }

  Float_t tofsigma = -999;

  const Bool_t ka = !(trk->GetStatus() & AliESDtrack::kTOFmismatch);
  const Bool_t kb =  (trk->GetStatus() & AliESDtrack::kTOFpid);
  const Bool_t ktof = ka && kb;
  if(ktof){
    tofsigma = fESDpid->NumberOfSigmasTOF( trk, itype, fESDpid->GetTOFResponse().GetTimeZero() );
  }

  return tofsigma;
}

void AliAnalysisTaskpkRaa::GetHMPID(const typetrack *trk, Float_t rawprob[], Float_t &hsig)
{
  if(!IsESD()){
    return;
  }

  Double_t tmpprob[5];
  for(Int_t ii=0; ii<5; ii++){
    rawprob[ii]=-999;
    tmpprob[ii]=-999;
  }

  hsig = -999;

  Float_t xmip, ymip;
  Int_t q, nphot;
  trk->GetHMPIDmip(xmip, ymip, q, nphot);

  Float_t xpc, ypc, theta, phi;
  trk->GetHMPIDtrk( xpc, ypc, theta, phi );

  if(q>120 && TMath::Sqrt((xpc-xmip)*(xpc-xmip)+(ypc-ymip)*(ypc-ymip)) < 1 ){
    trk->GetHMPIDpid(tmpprob);
    hsig = trk->GetHMPIDsignal();

    for(Int_t ii=0; ii<5; ii++){
      rawprob[ii]=tmpprob[ii];
    }
  }
}

Bool_t AliAnalysisTaskpkRaa::SetPDG(Int_t &outpdg, const Int_t input)
{
  //outpdg belongs to only 3 cases: not used, used once, used more than once = bad

  if(outpdg==fkBad)
    return kFALSE;

  if(outpdg!=fkNotUsed){
    outpdg = fkBad;
    return kFALSE;
  }

  outpdg=input;
  return kTRUE;
}

Int_t AliAnalysisTaskpkRaa::GetAlipdg(const Int_t rpdg)
{
  switch(TMath::Abs(rpdg)){
  case kElectron:
    return AliPID::kElectron;
  case kPiPlus:
    return AliPID::kPion;
  case kProton:
    return AliPID::kProton;
  default:
    //in GetEffMass "if(p1>4 || p2>4) return -1;"
    return 10;
  }
}

void AliAnalysisTaskpkRaa::GetV0type(const Int_t ntrk, const Float_t cen, Int_t pdgs[], Int_t mother[], Float_t cpa[], Float_t dca[], Float_t mass[], Float_t arm0[], Float_t arm1[], Float_t psi[])
{
  //cpa = CosineOfPointingAngle
  //ntrk all tracks

  if(!IsESD()){
    return;
  }

  for(Int_t ii=0; ii<ntrk; ii++){
    pdgs[ii]= fkNotUsed;
  }

  AliESDv0KineCuts v0cuts;
  //setevent will set primaryvertex automatically  
  v0cuts.SetEvent(fEvent);

  //for pp, cen all = -1
  v0cuts.SetMode(AliESDv0KineCuts::kPurity, cen>=0 ? AliESDv0KineCuts::kPbPb : AliESDv0KineCuts::kPP);

  for(Int_t iv0=0; iv0<fEvent->GetNumberOfV0s(); iv0++){
    AliESDv0 * esdv0 = fEvent->GetV0(iv0);
    if(!esdv0->GetOnFlyStatus())
      continue;

    Int_t pdgV0 = -999, pdgP = -999, pdgN = -999;
    if(v0cuts.ProcessV0(esdv0, pdgV0, pdgP, pdgN)){

      const Int_t itrkP = esdv0->GetPindex();
      const Int_t itrkN = esdv0->GetNindex();

      //if reused, pdgs[]==fkBad
      SetPDG(pdgs[itrkP], pdgP);
      SetPDG(pdgs[itrkN], pdgN);

      Float_t am[2];
      v0cuts.Armenteros(esdv0, am);

      //GetEffMass symmetric in p1, p2
      //only vaid using p1, p2, not the default!! can be seen for |mother|=3122

      //Float_t  GetV0CosineOfPointingAngle(Double_t refPointX, Double_t refPointY, Double_t refPointZ) const;
      //not really helpful, because for gamma and Lambda the decay vertex is not primary vertex, therefore refPoint is difficult to determine

      //set itrkP
      mother[itrkP] = pdgV0;
      cpa[itrkP] = esdv0->GetV0CosineOfPointingAngle();
      dca[itrkP] = esdv0->GetDcaV0Daughters();
      mass[itrkP] = esdv0->GetEffMass(GetAlipdg(pdgP), GetAlipdg(pdgN));
      arm0[itrkP] = am[0];
      arm1[itrkP] = am[1];
      psi[itrkP] = v0cuts.PsiPair(esdv0);

      //set itrkN
      mother[itrkN] = pdgV0;
      cpa[itrkN] = esdv0->GetV0CosineOfPointingAngle();
      dca[itrkN] = esdv0->GetDcaV0Daughters();
      mass[itrkN] = esdv0->GetEffMass(GetAlipdg(pdgP), GetAlipdg(pdgN));
      arm0[itrkN] = am[0];
      arm1[itrkN] = am[1];
      psi[itrkN] = v0cuts.PsiPair(esdv0);
    }
  }
}
